Post

UNet: Kiến trúc và Ứng Dụng trong Phân Đoạn Ảnh

Nội dung

1. Giới thiệu

UNet là kiến trúc mạng nơ-ron tích chập chuyên dùng cho bài toán phân đoạn ảnh, đặc biệt phổ biến trong y tế và viễn thám. Được giới thiệu bởi Olaf Ronneberger năm 2015, UNet nổi bật với kiến trúc đối xứng hình chữ “U”.



2. Kiến trúc UNet (Cập nhật)

UNet là mạng đối xứng hình chữ “U” gồm hai giai đoạn chính: downsampling (thu nhỏ không gian ảnh, tăng chiều sâu đặc trưng) và upsampling (phục hồi lại kích thước ảnh gốc). Điểm nổi bật của UNet là các skip connections kết nối trực tiếp các tầng tương ứng giữa hai giai đoạn, giúp bảo toàn thông tin không gian và cải thiện hiệu quả huấn luyện.

kiến trúc mô hình UNet kiến trúc mô hình UNet

  • Downsampling: Mỗi bước gồm hai lớp convolution (kernel 3x3, padding để giữ kích thước), kích hoạt ReLU, và lớp MaxPooling (2x2) giảm kích thước không gian đi một nửa. Qua mỗi bước, số kênh đặc trưng được tăng lên, giúp mạng học được các đặc trưng phức tạp hơn.

  • Upsampling: Dùng các lớp chuyển vị convolution (ConvTranspose) hoặc upsample để tăng kích thước ảnh lên gấp đôi, sau đó ghép (concatenate) với feature map tương ứng từ giai đoạn downsampling qua skip connection. Tiếp tục hai lớp convolution và ReLU để tinh chỉnh đặc trưng trước khi đưa sang bước tiếp theo.

  • Skip Connections: Giúp mạng giữ lại chi tiết không gian bị mất trong quá trình downsampling, từ đó cho kết quả phân đoạn chính xác hơn, đặc biệt với các chi tiết nhỏ hoặc biên phức tạp.

Giải thích kiến trúc U-Net

Nhìn vào hình trên, bạn sẽ thấy mô hình U-Net có hình chữ “U” đúng như tên gọi của nó. Mình sẽ chia nhỏ ra thành 3 phần chính để giải thích cho dễ hiểu, sau đó mình sẽ nói về luồng dữ liệu chi tiết


🔻 1. Downsampling – Làm ảnh nhỏ lại để lấy đặc trưng

Phần bên trái của mô hình là quá trình nén ảnh lại.

  • Ảnh đầu vào khá lớn, ví dụ 572x572x1.
  • Qua 2 lớp Conv 3x3 + ReLU → giữ gần nguyên kích thước nhưng số kênh tăng (vd: lên 64).
  • Sau đó là MaxPooling 2x2 (mũi tên đỏ) → giảm kích thước ảnh xuống còn một nửa, số kênh lại tiếp tục tăng.

Mỗi lần nén, ảnh nhỏ lại, nhưng chiều sâu (số kênh) lớn hơn → mô hình học được các đặc trưng quan trọng hơn.


🔺 2. Upsampling – Phóng to ảnh trở lại

Sau khi ảnh được nén đến mức trừu tượng nhất ở đáy chữ “U”, mô hình bắt đầu giải nén dần:

  • Up-conv 2x2 (mũi tên xanh lá): tăng kích thước ảnh gấp đôi.
  • Ghép (copy & crop) với ảnh từ bên trái cùng cấp độ (mũi tên xám) → bổ sung lại thông tin chi tiết đã mất.
  • Sau đó tiếp tục 2 lớp Conv 3x3 như bên trái.

Kết quả: ảnh dần được tái tạo lại, nhưng vẫn giữ những đặc trưng quan trọng.


3. Skip Connections – Cầu nối giữ chi tiết

Các mũi tên ngang màu xám chính là skip connection – đóng vai trò như “cầu nối”.

  • Chúng truyền thông tin từ phần nén (trái) sang phần giải nén (phải).
  • Giúp giữ lại những chi tiết quan trọng như cạnh, đường viền, biên đối tượng…

Nhờ đó, mô hình có thể vừa học đặc trưng tổng thể vừa giữ được độ chính xác ở mức pixel.


Luồng dữ liệu qua mô hình U-Net (dựa trên hình trên)

Giờ ta xem kỹ hơn luồng dữ liệu đi qua mô hình – từng bước ảnh được xử lý như thế nào.


Bước 1: Nhập ảnh

  • Ảnh đầu vào 572x572x1 (ảnh đen trắng).
  • Qua 2 lớp conv → giữ kích thước gần như nguyên, thành 568x568x64.

Bước 2: Downsampling – Nén dần

  • Pooling 2x2 → 284x284x64
  • Conv x2 → 280x280x128 → pooling → 140x140x128
  • Tiếp tục:
    138x138x256 → pooling
    68x68x512 → pooling
    → đáy: 28x28x1024

Tại đáy mô hình, ảnh rất nhỏ nhưng mang đặc trưng sâu nhất.


Bước 3: Upsampling – Phóng to dần

  • Upconv → 56x56x512
    → ghép với feature từ bên trái → conv x2 → 56x56x512
  • Tiếp tục:
    104x104x256
    200x200x128
    388x388x64

Bước 4: Kết thúc

  • Qua conv cuối cùng (1x1) → 388x388x2
    → mỗi pixel có 2 giá trị: xác suất thuộc lớp nào (foreground/background).

Tóm tắt kiến trúc U-Net

Giai đoạnKích thước ảnh (ví dụ)Mục đích chính
Nhập ảnh572x572x1Nhận ảnh đầu vào
Downsampling568→284→140→70→34→28Học đặc trưng trừu tượng, giảm kích thước
Bottleneck28x28x1024Đặc trưng sâu nhất
Upsampling28→56→104→200→388Tái tạo lại ảnh + kết hợp thông tin chi tiết
Output388x388x2Bản đồ phân đoạn (segmentation map)

Điểm mạnh của UNet là:

  • Cấu trúc đối xứng dễ hiểu.
  • Các kết nối skip giúp giữ chi tiết tốt.
  • Dễ huấn luyện và đạt hiệu quả cao trong phân đoạn ảnh y tế, vệ tinh,…

Thư viện hỗ trợ UNet

Segmentation Models PyTorch (SMP) là một thư viện Python rất tiện lợi để xây dựng và huấn luyện các mô hình phân đoạn ảnh, trong đó có UNet. SMP cung cấp các backbone (phần CNN trích xuất đặc trưng) đa dạng, giúp cải thiện hiệu quả mô hình một cách dễ dàng.

Ví dụ đơn giản sử dụng UNet với backbone ResNet:

1
2
3
4
5
6
7
8
import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # backbone CNN
    encoder_weights="imagenet",     # sử dụng trọng số học sẵn
    in_channels=3,                  # số kênh ảnh đầu vào (RGB)
    classes=1                       # số lớp đầu ra (ví dụ: 1 lớp cho phân đoạn nhị phân)
)

3. Ứng dụng UNet

UNet được dùng nhiều trong các bài toán segmentation:

  • Ảnh y tế: phân vùng mô, cơ quan
  • Viễn thám: phân mảnh đất, sông, nhà
  • Đếm thực thể: tế bào, hạt giống…

4. Ưu điểm và hạn chế

Ưu điểm:

  • Hoạt động tốt với ít dữ liệu
  • Tái tạo chi tiết nhờ skip connections
  • Kiến trúc đơn giản, dễ mở rộng (UNet++, ResUNet,…)

Hạn chế:

  • Kích thước ảnh cố định (nếu không padding phù hợp)
  • Có thể nặng RAM với ảnh lớn

5. Tham khảo

[1] Scikit-learn Developers. sklearn.metrics.confusion_matrix — scikit-learn 1.4.2 documentation. 🔗 https://scikit-learn.org/stable/modules/generated/sklearn.metrics.confusion_matrix.html
[2] Segmentation Models PyTorch. UNet model documentation. 🔗 https://smp.readthedocs.io/en/latest/models.html#unet

This post is licensed under CC BY 4.0 by the author.